""" 
    本程序通过调用Google OR-Tools工具包求解线性规划模型，实现基于列车运行图压缩法的铁路区间通过能力计算
    作者：廖正文 北京交通大学

    运行本程序前，需要安装Google OR-Tools的Python包，安装方法如下：
    保持计算机在可以访问互联网的状态，在命令提示符中，输入
    python -m pip install --user ortools
    即可完成安装

    如果安装过程较慢，可以尝试使用清华大学TUNA的镜像，即输入
    python -m pip install --user ortools -i https://pypi.tuna.tsinghua.edu.cn/simple
 """

from functools import cmp_to_key
from ortools.linear_solver import pywraplp  # 载入 Google OR-Tools 求解工具包
import PyTimetableCompressionData as ptcdata
import os


### 全局变量 ###
Data = None
''' 数据 '''

# appliedSolverName = "Gurobi"  # 设置Gurobi作为线性规划的求解器
appliedSolverName = "SCIP"  # 设置SCIP作为线性规划的求解器
# appliedSolverName = "CBC"  # 设置CBC作为线性规划的求解器
# appliedSolverName = "GLOP"  # 设置GLOP作为线性规划的求解器
# appliedSolverName = "CP-SAT"  # 设置CP-SAT作为线性规划的求解器


''' 采用求解器名称 '''

Solver = pywraplp.Solver.CreateSolver(appliedSolverName)
''' 求解器 '''

trainSequenceDict = dict()
''' 列车进入区间顺序字典 '''

VarDict = dict()
''' 决策变量字典 '''

### 全局变量 ###

### 列车排序相关 ###
sortingSegment = None
''' 当前排序的区间 '''


def CompareEnteringSequence(train1: ptcdata.Train, train2: ptcdata.Train):
    ''' 对比两列车进入区间的时刻，train1先进入返回-1，train2先进入返回1 '''

    # 按原图的顺序安排越行
    train1EnteringTime = [event for event in train1.OriginalTrainEventList
                          if event.Type == "Departure" and event.Station == sortingSegment.FromStation][0].Time
    train2EnteringTime = [event for event in train2.OriginalTrainEventList
                          if event.Type == "Departure" and event.Station == sortingSegment.FromStation][0].Time

    # 不安排列车越行
    # train1EnteringTime = train1.OriginalTrainEventList[0].Time
    # train2EnteringTime = train2.OriginalTrainEventList[0].Time

    if (train1EnteringTime < train2EnteringTime):
        return -1
    else:
        return 1


def SortTrainListBySegment(trainList, segment):
    ''' 根据原始列车时刻表，按列车进入区间的顺序对列车列表进行排序 '''
    global sortingSegment
    sortingSegment = segment
    trainList.sort(key=cmp_to_key(CompareEnteringSequence))
    return trainList


def Initialize():
    ''' 初始化 '''

    for segment in Data.SegmentList:
        trainList = [
            train for train in Data.TrainList if segment in train.RouteSegmentList]
        sortedTrainList = SortTrainListBySegment(trainList, segment)
        trainSequenceDict[segment] = sortedTrainList

    if (os.path.exists("OutputFiles") != 1):  # 创建输出文件目录
        os.mkdir("./OutputFiles")

    OutputTrainSequenceBySegment()  # 输出列车进入区间顺序数据


def OutputTrainSequenceBySegment():
    ''' 输出列车进入区间顺序数据 '''

    f = open("./OutputFiles/Debug_TrainSequence.csv", "w")
    f.writelines("FromStation,ToStation,TrainID,Seq\n")
    for segment in Data.SegmentList:
        trainList = trainSequenceDict[segment]
        for train in trainList:
            f.writelines("{0},{1},{2},{3}\n".format(
                segment.FromStation.Name, segment.ToStation.Name, train.ID, trainList.index(train)))
    f.close()

### 列车排序相关 ###


def BuildModel():
    """ 创建运行图压缩线性规划模型 """

    BuildVariables()  # 构建决策变量
    BuildConstraints()  # 构建约束条件
    BuildObjective()  # 构建目标函数


def BuildVariables():
    ''' 创建决策变量 '''

    for train in Data.TrainList:
        for station in train.RouteStationList:
            a_VarName = "a_{0}_{1}".format(train.ID, station.ID)
            # 列车 train 在车站 station 的到达时刻（模型中的变量 a_f^s）
            a = Solver.IntVar(Data.TimetableStartTime,
                              Data.TimetableEndTime, a_VarName)
            VarDict[a_VarName] = a

            d_VarName = "d_{0}_{1}".format(train.ID, station.ID)
            # 列车 train 在车站 station 的出发时刻（模型中的变量 d_f^s）
            d = Solver.IntVar(Data.TimetableStartTime,
                              Data.TimetableEndTime, d_VarName)
            VarDict[d_VarName] = d


def BuildObjective():
    ''' 创建目标函数 '''

    objExpr1 = 0
    for train in Data.TrainList:
        lastStation = train.RouteStationList[train.RouteStationList.__len__(
        ) - 1]
        d = VarDict["d_{0}_{1}".format(train.ID, lastStation.ID)]
        objExpr1 += d

    objExpr2 = 0
    for train in Data.TrainList:
        firstStation = train.RouteStationList[0]
        lastStation = train.RouteStationList[train.RouteStationList.__len__(
        ) - 1]
        a = VarDict["a_{0}_{1}".format(train.ID, firstStation.ID)]
        d = VarDict["d_{0}_{1}".format(train.ID, lastStation.ID)]
        objExpr2 += d - a

    Solver.Minimize(objExpr1 + 0.1 * objExpr2)  # 目标函数为最小化所有列车到发时刻之和


def BuildConstraints():
    ''' 创建约束条件 '''

    BuildRunningTimeConstraints()  # 创建区间运行时分约束
    BuildMinDwellTimeConstraint()  # 创建最小停站时分约束
    BuildHeadwayConstraint()  # 创建追踪间隔时间约束


def BuildRunningTimeConstraints():
    ''' 创建列车区间运行时分约束条件 '''

    for train in Data.TrainList:  # 对于每列车
        for segment in train.RouteSegmentList:  # 对于列车运行径路上的每个区间
            d_fromStation = VarDict["d_{0}_{1}".format(
                train.ID, segment.FromStation.ID)]  # 区间出发车站的出发时刻
            a_toStation = VarDict["a_{0}_{1}".format(
                train.ID, segment.ToStation.ID)]  # 区间到达车站的到达时刻

            fromStationStopTime = train.GetStopTime(segment.FromStation)
            toStationStopTime = train.GetStopTime(segment.ToStation)

            if fromStationStopTime > 0:
                isAccelerate = True
            else:
                isAccelerate = False

            if toStationStopTime > 0:
                isDecelerate = True
            else:
                isDecelerate = False

            runningTime = segment.GetRunningTime(
                train.Type, isAccelerate, isDecelerate)  # 区间运行时分
            Solver.Add(a_toStation - d_fromStation ==
                       runningTime)  # 生成区间运行时分约束


def BuildMinDwellTimeConstraint():
    ''' 创建列车最小停站时分约束条件 '''

    for train in Data.TrainList:  # 对于每列车
        for station in train.RouteStationList:  # 对于列车运行径路上的每个车站
            a = VarDict["a_{0}_{1}".format(train.ID, station.ID)]  # 列车在车站的到达时刻
            d = VarDict["d_{0}_{1}".format(train.ID, station.ID)]  # 列车在车站的出发时刻

            planStopTime = train.GetStopTime(station)

            if planStopTime == 0:
                Solver.Add(d - a == 0)  # 生成列车在车站的最小停站时分约束（列车不停站通过）
            else:
                minDwellTime = 2  # 列车在车站的最小停站时分
                Solver.Add(d - a >= minDwellTime)  # 生成列车在车站的最小停站时分约束（列车停站）


def BuildHeadwayConstraint():
    ''' 创建列车最小追踪间隔时间约束条件 '''

    for segment in Data.SegmentList:
        trainList = trainSequenceDict[segment]
        for i in range(1, trainList.__len__()):
            leadingTrain = trainList[i - 1]
            followingTrain = trainList[i]

            d_leading = VarDict["d_{0}_{1}".format(
                leadingTrain.ID, segment.FromStation.ID)]
            d_following = VarDict["d_{0}_{1}".format(
                followingTrain.ID, segment.FromStation.ID)]

            leadingTrainFromStationStop = leadingTrain.GetStopTime(
                segment.FromStation)
            followingTrainFromStationStop = followingTrain.GetStopTime(
                segment.FromStation)

            headway = segment.GetHeadway("0", "Departure",
                                         True if leadingTrainFromStationStop > 0 else False,
                                         True if followingTrainFromStationStop > 0 else False)

            minDepartureHeadway = headway

            Solver.Add(d_following - d_leading >= minDepartureHeadway)

            a_leading = VarDict["a_{0}_{1}".format(
                leadingTrain.ID, segment.ToStation.ID)]
            a_following = VarDict["a_{0}_{1}".format(
                followingTrain.ID, segment.ToStation.ID)]

            leadingTrainToStationStop = leadingTrain.GetStopTime(
                segment.ToStation)
            followingTrainToStationStop = followingTrain.GetStopTime(
                segment.ToStation)

            headway = segment.GetHeadway("0", "Arrival",
                                         True if leadingTrainToStationStop > 0 else False,
                                         True if followingTrainToStationStop > 0 else False)

            minArrivalHeadway = headway
            Solver.Add(a_following - a_leading >= minArrivalHeadway)


def ParseSolution():
    ''' 自求解器中解析求解结果，生成压缩后的列车运行图数据 '''

    for train in Data.TrainList:
        for station in train.RouteStationList:
            a = VarDict["a_{0}_{1}".format(train.ID, station.ID)]  # 列车在车站的到达时刻
            d = VarDict["d_{0}_{1}".format(train.ID, station.ID)]  # 列车在车站的出发时刻

            arrivalTime = a.SolutionValue()
            departureTime = d.SolutionValue()

            arrivalEvent = ptcdata.TrainEvent()
            arrivalEvent.Station = station
            arrivalEvent.Train = train
            arrivalEvent.Type = "Arrival"
            arrivalEvent.Time = int(arrivalTime)
            train.CompressedTrainEventList.append(arrivalEvent)

            departureEvent = ptcdata.TrainEvent()
            departureEvent.Station = station
            departureEvent.Train = train
            departureEvent.Type = "Departure"
            departureEvent.Time = int(departureTime)
            train.CompressedTrainEventList.append(departureEvent)

    print("求解数据解析完成！")


def WriteLpFile(fileName):
    ''' 
        输出*.lp文件字符串，并保存至指定文件中
        lp格式的文件是线性（整数）规划模型完全展开后的形式，可以根据此文件检查以上生成模型的代码是否存在问题
    '''
    lp_str = Solver.ExportModelAsLpFormat(False)
    lp_file = open(fileName, "w")
    lp_file.writelines(lp_str)
    lp_file.close()


def Solve():
    """ 求解运行图压缩线性规划模型 """

    Solver.EnableOutput()  # 允许求解引擎向终端输出日志
    WriteLpFile("ResourceBasedModel.lp")  # 输出模型lp文件

    if Solver.Solve() == pywraplp.Solver.OPTIMAL:
        ParseSolution()
        print("求解结束！")
        return True

    return False
